from ultralytics import YOLO
import numpy as np
import cv2
# Load a model


from PIL import Image

def get_center(point):
    center_x = round(sum([p[0] for p in point])/len(point),4)
    center_y = round(sum([p[1] for p in point]) / len(point), 4)
    return (center_x,center_y)


def calculate_distance( points, point_center):
    infos =  []
    for point in points.tolist():
        pixel_distance = np.sqrt((point_center[0] - point[0]) ** 2 + (point_center[1] - point[1]) ** 2)
        infos.append((pixel_distance,point))
    infos.sort(key=lambda x:x[0],reverse=True)
    return (int((infos[0][1][0] + infos[1][1][0])/2),int((infos[0][1][1] + infos[1][1][1])/2))



def get_start_point(points, point_center,stard_point):
    infos = []
    for point in points.tolist():
        pixel_distance = np.sqrt((point_center[0] - point[0]) ** 2 + (point_center[1] - point[1]) ** 2)
        infos.append((pixel_distance, point))
    infos.sort(key=lambda x: x[0], reverse=True)

    infos = infos[:2]
    final_point = None
    final_value = 10000000
    for _,point in infos:
        pixel_distance = np.sqrt((stard_point[0] - point[0]) ** 2 + (stard_point[1] - point[1]) ** 2)
        if pixel_distance < final_value:
            final_value = pixel_distance
            final_point = point

    return final_point


def get_point_info(obb):
    info = {}
    for i, point,stard in zip(obb.cls.cpu(), obb.xyxyxyxy.cpu().numpy(),obb.xyxy.cpu().numpy()):
        if i == 0 and 'center' not in info:
            info['center'] = get_center(point)
            info['center_origin'] = point

    for i, point,stard in zip(obb.cls.cpu(), obb.xyxyxyxy.cpu().numpy(),obb.xyxy.cpu().numpy()):
        if i == 1 and 'point' not in info:
            info["point"] = calculate_distance(point,info['center'])
            info['point_origin'] = point
        elif i == 2 and 'end' not in info:
            info["end"] = calculate_distance(point,info['center'])
            info['end_origin'] = point
        elif i == 3 and 'start' not in info:
            info["start"] = get_start_point(point,info['center'],stard)
            info['start_origin'] = point
            info["start_stard"] = stard
    return info


def GetClockAngle( v1, v2):
    # 2个向量模的乘积 ,返回夹角
    TheNorm = np.linalg.norm(v1) * np.linalg.norm(v2)
    # 叉乘
    rho = np.rad2deg(np.arcsin(np.cross(v1, v2) / TheNorm))
    # 点乘
    theta = np.rad2deg(np.arccos(np.dot(v1, v2) / TheNorm))
    if rho > 0:
        return theta
    else:
        return 360 - theta


def get_theho_result(info,max_value,zero_offset):
    start_x,start_y = info["start"]
    end_x,end_y = info["end"]
    center_x,center_y = info["center"]
    point_x,point_y = info["point"]

    v1 = [start_x - center_x,start_y - center_y]
    v2 = [point_x - center_x,point_y - center_y]
    theta = GetClockAngle(v1, v2)
    start_point_theta = round(theta, 3)

    v3 = [start_x - center_x, start_y - center_y]
    v4 = [end_x - center_x, end_y - center_y]

    theta2 = GetClockAngle(v3, v4)
    start_end_theta = round(theta2, 3)

    value = round(round((max_value-zero_offset) * start_point_theta/start_end_theta ,3) + zero_offset,2)
    return value


if __name__ == '__main__':
    model = YOLO(
        r"F:\yaoyao\images_project\energy-meter\obb_train\runs\obb\train6\weights\best.pt")  # pretrained YOLO11n model

    images  = [
        r"E:\github\yolov8-gradio\samples\1.jpg",
        #r"F:\yaoyao\images_project\energy-meter\data\final_train\2.jpg",
        #r"F:\yaoyao\images_project\energy-meter\data\final_train\3.jpg",
        ##r"F:\yaoyao\images_project\energy-meter\data\final_train\4.jpg"
    ]



    for image in images:
        # Run batched inference on a list of images
        results = model.predict([image]
                                ,conf=0.01,iou=0.5)  # return a list of Results objects

        # Process results list
        for result in results:
            probs = result.probs  # Probs object for classification outputs
            obb = result.obb  # Oriented boxes object for OBB outputs

            info = get_point_info(obb)
            #result.names = {0: '', 1: '', 2: '', 3: ''}
            j = 1
            get_theho_result(info,60,1)

            frame = cv2.imread(image)

            point1 = [int(i) for i in info["start"]]
            point2 = [int(i) for i in info["center"]]
            point3 = [int(i) for i in info["end"]]
            point4 = [int(i) for i in info["point"]]
            cv2.line(frame, point1, point2, (0, 255, 0), 2)
            cv2.line(frame, point3, point2, (0, 255, 0), 2)
            cv2.line(frame, point4, point2, (0, 255, 0), 2)

            point6 = [int(i) for i in info["start_stard"]]
            #cv2.line(frame, point5, point6, (0, 255, 0), 2)
            cv2.rectangle(frame,point6, (0, 255, 0), 2)


            point5 = [int(i) for i in info["start_origin"][0]]
            point6 = [int(i) for i in info["start_origin"][1]]
            point7 = [int(i) for i in info["start_origin"][3]]
            cv2.line(frame, point5, point6, (0, 255, 0), 2)
            cv2.line(frame, point5, point7, (0, 255, 0), 2)
            cv2.imshow('4-Hao-Zhuan-Lu-Ce-Shi', frame)
            cv2.waitKey(0)

            result.show()  # display to screen
            result.save(filename="result.jpg")  # save to disk










